import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

class RobotArmVisualizer:
    def __init__(self):
        # 初始参数设置
        self.L1 = 80         # 连杆1长度
        self.L2 = 60         # 连杆2长度
        self.L3 = 40         # 连杆3长度
        
        # 初始末端位置（在工作空间内）
        self.target_x = 100
        self.target_y = 50
        
        # 通过逆运动学计算初始关节角度（下肘位）
        self.theta1, self.theta2 = self.inverse_kinematics(self.target_x, self.target_y)
        self.theta3 = 270 - self.theta1 - self.theta2  # 保持角度约束
        
        # 创建图形和轴
        self.fig, self.ax = plt.subplots(figsize=(10, 10))
        plt.subplots_adjust(left=0.1, bottom=0.5)
        
        # 绘制运动范围散点（只显示下肘位）
        self.workspace_points, = self.ax.plot([], [], 'b.', alpha=0.2, label='下肘位运动范围')
        
        # 绘制机械臂
        self.arm_line, = self.ax.plot([], [], 'o-', lw=5, color='purple', 
                                    markersize=10, markerfacecolor='red')
        self.target_marker, = self.ax.plot([], [], 'rx', markersize=15, label='目标位置')
        self.end_effector_text = self.ax.text(0, 0, '', fontsize=12)
        
        # 标记目标工作空间
        self.ax.axvspan(60, 150, ymin=0, ymax=1, facecolor='yellow', alpha=0.2)
        self.ax.text(105, 120, '目标工作空间', ha='center', fontsize=10, 
                    bbox=dict(facecolor='yellow', alpha=0.5))
        
        # 绘制边界线
        self.ax.axvline(x=60, color='red', linestyle='--', alpha=0.7, label='X=60')
        self.ax.axvline(x=150, color='red', linestyle='--', alpha=0.7, label='X=150')
        self.ax.axhline(y=0, color='red', linestyle='--', alpha=0.7, label='Y=0')
        
        # 设置坐标轴
        self.ax.set_xlim(0, 250)
        self.ax.set_ylim(0, 250)
        self.ax.set_xlabel('X坐标')
        self.ax.set_ylabel('Y坐标')
        self.ax.set_title('三连杆机械臂（下肘位约束 + 逆运动学）')
        self.ax.grid(True)
        self.ax.legend()
        self.ax.set_aspect('equal', adjustable='box')
        
        # 创建滑块区域
        slider_color = 'lightgoldenrodyellow'
        
        # 末端位置滑块（X和Y）
        ax_target_x = plt.axes([0.1, 0.40, 0.8, 0.03], facecolor=slider_color)
        ax_target_y = plt.axes([0.1, 0.35, 0.8, 0.03], facecolor=slider_color)
        
        # 关节角度显示
        self.joint_text = self.fig.text(0.5, 0.32, 
                                      f'θ1 = {self.theta1:.1f}°, θ2 = {self.theta2:.1f}°, θ3 = {self.theta3:.1f}°\n' +
                                      f'约束: θ1+θ2+θ3 = {self.theta1+self.theta2+self.theta3:.1f}°', 
                                      ha='center', fontsize=10,
                                      bbox=dict(facecolor='lightblue', alpha=0.5))
        
        # 连杆长度滑块
        ax_L1 = plt.axes([0.1, 0.25, 0.8, 0.03], facecolor=slider_color)
        ax_L2 = plt.axes([0.1, 0.20, 0.8, 0.03], facecolor=slider_color)
        ax_L3 = plt.axes([0.1, 0.15, 0.8, 0.03], facecolor=slider_color)
        
        # 创建滑块控件
        self.slider_target_x = Slider(ax_target_x, '目标X坐标', 60, 150, valinit=self.target_x)
        self.slider_target_y = Slider(ax_target_y, '目标Y坐标', 0, 200, valinit=self.target_y)
        self.slider_L1 = Slider(ax_L1, r'$L_1$', 50, 120, valinit=self.L1)
        self.slider_L2 = Slider(ax_L2, r'$L_2$', 40, 100, valinit=self.L2)
        self.slider_L3 = Slider(ax_L3, r'$L_3$', 20, 80, valinit=self.L3)
        
        # 注册更新函数
        self.slider_target_x.on_changed(self.update_target_position)
        self.slider_target_y.on_changed(self.update_target_position)
        self.slider_L1.on_changed(self.update_link_lengths)
        self.slider_L2.on_changed(self.update_link_lengths)
        self.slider_L3.on_changed(self.update_link_lengths)
        
        # 初始化显示
        self.update_workspace()
        self.update_arm()
        
    def inverse_kinematics(self, x, y):
        """
        逆运动学求解，只返回下肘位（elbow down）解
        下肘位定义为theta2为负值（连杆2向下弯曲）
        """
        # 考虑末端执行器长度的偏移（由于theta3固定为270-theta1-theta2）
        offset_x = self.L3 * np.cos(np.radians(270))  # cos(270°) = 0
        offset_y = self.L3 * np.sin(np.radians(270))  # sin(270°) = -1
        
        # 计算腕点（关节3）位置
        wrist_x = x - offset_x
        wrist_y = y - offset_y
        
        # 计算从基座到腕点的距离
        D = np.sqrt(wrist_x**2 + wrist_y**2)
        
        # 检查可达性
        if D > self.L1 + self.L2 or D < abs(self.L1 - self.L2):
            # 如果目标不可达，返回最近可达位置的解
            D = np.clip(D, abs(self.L1 - self.L2), self.L1 + self.L2)
        
        # 计算theta1（基座旋转角）
        alpha = np.degrees(np.arctan2(wrist_y, wrist_x))
        beta = np.degrees(np.arccos((self.L1**2 + D**2 - self.L2**2) / (2 * self.L1 * D)))
        
        # 下肘位解（theta2为负）
        theta1 = alpha + beta
        gamma = np.degrees(np.arccos((self.L1**2 + self.L2**2 - D**2) / (2 * self.L1 * self.L2)))
        theta2 = -(180 - gamma)  # 下肘位解（与上肘位相反）
        
        # 确保角度在合理范围内
        theta1 = (theta1 + 360) % 360
        if theta1 > 180:
            theta1 -= 360
            
        theta2 = np.clip(theta2, -170, 170)  # 限制关节2的运动范围
        
        return theta1, theta2
    
    def forward_kinematics(self, theta1, theta2):
        """正运动学计算，返回各关节位置"""
        theta3 = 270 - theta1 - theta2
        
        # 转换为弧度
        t1 = np.radians(theta1)
        t2 = np.radians(theta2)
        t3 = np.radians(theta3)
        
        # 计算各关节绝对角度
        angle1 = t1
        angle2 = t1 + t2
        angle3 = t1 + t2 + t3  # 应为270度
        
        # 各关节坐标
        x0, y0 = 0, 0
        x1 = x0 + self.L1 * np.cos(angle1)
        y1 = y0 + self.L1 * np.sin(angle1)
        x2 = x1 + self.L2 * np.cos(angle2)
        y2 = y1 + self.L2 * np.sin(angle2)
        x3 = x2 + self.L3 * np.cos(angle3)
        y3 = y2 + self.L3 * np.sin(angle3)
        
        return [(x0, y0), (x1, y1), (x2, y2), (x3, y3)]
    
    def generate_workspace(self):
        """生成下肘位的运动范围"""
        x_range = np.arange(60, 151, 5)
        y_range = np.arange(0, 201, 5)
        
        points = []
        for x in x_range:
            for y in y_range:
                # 计算逆运动学解
                theta1, theta2 = self.inverse_kinematics(x, y)
                # 只保留下肘位（theta2 < 0）
                if theta2 < 0:
                    # 验证正运动学是否匹配（考虑计算误差）
                    joints = self.forward_kinematics(theta1, theta2)
                    end_x, end_y = joints[-1]
                    if abs(end_x - x) < 5 and abs(end_y - y) < 5:
                        points.append((x, y))
        
        return np.array(points)
    
    def update_workspace(self):
        """更新下肘位运动范围"""
        workspace = self.generate_workspace()
        if len(workspace) > 0:
            self.workspace_points.set_xdata(workspace[:, 0])
            self.workspace_points.set_ydata(workspace[:, 1])
    
    def update_target_position(self, val):
        """更新目标位置并计算对应的关节角度（下肘位）"""
        self.target_x = self.slider_target_x.val
        self.target_y = self.slider_target_y.val
        
        # 计算下肘位的逆运动学解
        self.theta1, self.theta2 = self.inverse_kinematics(self.target_x, self.target_y)
        self.theta3 = 270 - self.theta1 - self.theta2
        
        self.update_arm()
    
    def update_link_lengths(self, val):
        """更新连杆长度并重新计算工作空间"""
        self.L1 = self.slider_L1.val
        self.L2 = self.slider_L2.val
        self.L3 = self.slider_L3.val
        
        # 更新工作空间
        self.update_workspace()
        
        # 重新计算当前目标位置的逆解
        self.theta1, self.theta2 = self.inverse_kinematics(self.target_x, self.target_y)
        self.theta3 = 270 - self.theta1 - self.theta2
        
        self.update_arm()
    
    def update_arm(self):
        """更新机械臂显示"""
        # 计算各关节位置
        positions = self.forward_kinematics(self.theta1, self.theta2)
        x = [p[0] for p in positions]
        y = [p[1] for p in positions]
        
        # 更新机械臂显示
        self.arm_line.set_xdata(x)
        self.arm_line.set_ydata(y)
        
        # 更新目标位置标记
        self.target_marker.set_xdata([self.target_x])
        self.target_marker.set_ydata([self.target_y])
        
        # 更新末端位置文本
        end_x, end_y = positions[-1]
        self.end_effector_text.set_text(f'末端位置: ({end_x:.2f}, {end_y:.2f})')
        self.end_effector_text.set_x(end_x)
        self.end_effector_text.set_y(end_y)
        
        # 更新关节角度显示
        self.joint_text.set_text(
            f'θ1 = {self.theta1:.1f}°, θ2 = {self.theta2:.1f}°, θ3 = {self.theta3:.1f}°\n' +
            f'约束: θ1+θ2+θ3 = {self.theta1+self.theta2+self.theta3:.1f}°'
        )
        
        self.fig.canvas.draw_idle()
    
    def show(self):
        plt.show()

if __name__ == "__main__":
    visualizer = RobotArmVisualizer()
    visualizer.show()
